package com.zhiyou.app.common.ws;


import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.zhiyou.app.common.util.SseEmitterSender;
import com.zhiyou.app.pojo.entity.OrderEntity;
import com.zhiyou.app.pojo.entity.TaskEntity;
import com.zhiyou.app.pojo.enums.OrderStatusEnum;
import com.zhiyou.app.pojo.enums.TaskStatusEnum;
import com.zhiyou.app.service.ComfyUIService;
import com.zhiyou.app.service.IOrderService;
import com.zhiyou.app.service.ITaskService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.Resource;
import java.util.List;
import java.util.Map;

@Slf4j
@Component
public class DefaultTaskFinishHandler implements TaskFinishHandler {
    @Resource
    private ComfyUIService comfyUIService;
    @Resource
    private ITaskService taskService;
    @Resource
    private IOrderService orderService;

    @Transactional(rollbackFor = Exception.class)
    @Override
    public void handle(String promptId, String orderNum, TaskEntity taskEntity, ComfyuiTaskInfo taskValue) {
        if (taskValue.getStatus().equals("error")) {
            SseEmitterSender.send(orderNum, "[state]:" + "任务异常");
            log.info("任务[" + taskValue.getTask_id() + "]异常:" + JSONUtil.toJsonStr(taskValue.getError_info()));
            taskEntity.setTaskError(JSONUtil.toJsonStr(taskValue.getError_info()));
            taskEntity.setTaskStatus(TaskStatusEnum.ERROR);
            taskService.updateById(taskEntity);
            orderService.update(new LambdaUpdateWrapper<OrderEntity>().eq(OrderEntity::getOrderNum, taskEntity.getOrderNum())
                    .set(OrderEntity::getOrderStatus, OrderStatusEnum.CLOSED).set(OrderEntity::getOrderDesc, "任务执行结果异常，关闭订单")
            );
        } else {
            Map<String, List<String>> res = comfyUIService.getRes(promptId, taskEntity);
            log.info("onFinished:{},res:{}", promptId, res);
            taskEntity.setTaskResult(JSONUtil.toJsonStr(res));
            if (res == null || res.isEmpty()) {
                taskEntity.setTaskStatus(TaskStatusEnum.ERROR);
                taskEntity.setTaskDesc("任务执行结果异常");
            } else {
                taskEntity.setTaskStatus(TaskStatusEnum.COMPLETE);
            }
            taskService.updateById(taskEntity);
            log.info("任务完成更新订单状态：{}", promptId);
            orderService.update(new LambdaUpdateWrapper<OrderEntity>().eq(OrderEntity::getOrderNum, taskEntity.getOrderNum())
                    .set(OrderEntity::getOrderStatus, OrderStatusEnum.FINISHED)
            );
            SseEmitterSender.send(orderNum, "[result]:" + taskService.getTaskResult(taskEntity));
        }
    }
}
